# -*- coding:utf-8 -*-

import argparse
import os
from tqdm import tqdm

from funasr import AutoModel

# 固定模型路径
path_asr = '/mnt/workspace/auto_biaozhu/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch'
path_vad = '/mnt/workspace/auto_biaozhu/speech_fsmn_vad_zh-cn-16k-common-pytorch'
path_punc = '/mnt/workspace/auto_biaozhu/punc_ct-transformer_zh-cn-common-vocab272727-pytorch'

# 加载模型
model = AutoModel(
    model=path_asr,
    model_revision="v2.0.4",
    vad_model=path_vad,
    vad_model_revision="v2.0.4",
    punc_model=path_punc,
    punc_model_revision="v2.0.4",
)

def process_audio_files(input_folder, output_file_path):
    count = 0
    with open(output_file_path, "w", encoding="utf-8") as output_file:
        for root, dirs, files in os.walk(input_folder):
            for name in tqdm(files):
                if name.lower().endswith('.wav'):
                    input_file_path = os.path.join(root, name)
                    folder_name = os.path.basename(root)
                    try:
                        text = model.generate(input=input_file_path)[0]["text"]
                        output_line = f"{input_file_path}|{folder_name}|ZH|{text}\n"
                        output_file.write(output_line)
                        count += 1
                    except Exception as e:
                        print(f"Error processing {name}: {str(e)}")
    print(f"总共转写音频数量为: {count}")

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("-i", "--input_folder", type=str, required=True,
                        help="Path to the folder containing WAV files.")
    parser.add_argument("-o", "--output_file", type=str, required=True,
                        help="Output file path for the transcription list.")
    args = parser.parse_args()
    process_audio_files(args.input_folder, args.output_file)